package com.bsfcs.rag.repository;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Repository;
import org.springframework.util.CollectionUtils;

import java.util.*;

@Repository
@Primary
public class TenantAwareChatMemoryRepository implements ChatMemoryRepository {

    private final JdbcTemplate jdbcTemplate;

    public TenantAwareChatMemoryRepository(JdbcTemplate jdbcTemplate) {
        this.jdbcTemplate = jdbcTemplate;
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        String[] params = conversationId.split(":");
        return jdbcTemplate.query(
                "SELECT role, content FROM chat_messages WHERE conversation_id = ? ORDER BY created_at ASC",
                (rs, rowNum) -> {
                    String role = rs.getString("role");
                    String content = rs.getString("content");

                    switch (role.toUpperCase()) {
                        case "USER":
                            return new org.springframework.ai.chat.messages.UserMessage(content);
                        case "ASSISTANT":
                            return new org.springframework.ai.chat.messages.AssistantMessage(content);
                        case "SYSTEM":
                            return new org.springframework.ai.chat.messages.SystemMessage(content);
                        default:
                            throw new IllegalArgumentException("未知的消息角色: " + role);
                    }
                },
                params[0]
        );
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        if (messages.isEmpty()) {
            return;
        }

        String[] params = conversationId.split(":");

        // 保存会话信息到 chat_conversations 表
        String conversationCheckSql = "SELECT COUNT(1) FROM chat_conversations WHERE id = ?";
        Integer count = jdbcTemplate.queryForObject(conversationCheckSql, Integer.class, params[0]);

        if (count != null && count == 0) {
            String conversationSql = "INSERT INTO chat_conversations (id, tenant_id, user_id, title, created_at) VALUES (?, ?, ?, ?, now())";

            String title = "对话-" + messages.get(0).getText().substring(0, Math.min(messages.get(0).getText().length(), 8));
            jdbcTemplate.update(conversationSql, params[0], params[2], params[1], title);
        }


        // 添加去重逻辑：避免重复保存相同的消息
        Map<String, Message> messageMap = new HashMap<>();
        for (Message msg : messages) {
            String originalContent = msg.getText();
            String actualQuestion = "";
            // 提取"问题:"后面的内容作为实际问题
            if ("user".equals(msg.getMessageType().getValue()) && originalContent.contains("用户问题:")) {
                actualQuestion = originalContent.substring(originalContent.indexOf("用户问题:") + 5).trim();
            } else {
                actualQuestion = originalContent;
            }
            String checkExistingSql = "SELECT id FROM chat_messages WHERE conversation_id = ? AND role = ? AND content = ?";
            List<String> existing = jdbcTemplate.queryForList(checkExistingSql, String.class, params[0], msg.getMessageType().getValue(), actualQuestion);
            if (CollectionUtils.isEmpty(existing)) {
                messageMap.put(actualQuestion, msg);
            }
        }


        String sql = "INSERT INTO chat_messages (conversation_id, role, content, tenant_id, user_id, id, created_at) " +
                "VALUES (?, ?, ?, ?, ?, ?, now())";

        jdbcTemplate.batchUpdate(sql, messageMap.values(), messageMap.size(), (ps, msg) -> {
            String role = msg.getMessageType().getValue();

            String originalContent = msg.getText();
            String actualQuestion = "";
            // 提取"问题:"后面的内容作为实际问题
            if ("user".equals(role) && originalContent.contains("用户问题:")) {
                actualQuestion = originalContent.substring(originalContent.indexOf("用户问题:") + 5).trim();
            } else {
                actualQuestion = originalContent;
            }

            ps.setString(1, params[0]);
            ps.setString(2, role);
            ps.setString(3, actualQuestion);
            ps.setString(4, params[2]);
            ps.setString(5, params[1]);
            ps.setString(6, UUID.randomUUID().toString());
        });
    }

    // 新增方法：更新会话标题
    public void updateConversationTitle(String conversationId, String title) {
        String sql = "UPDATE chat_conversations SET title = ? WHERE id = ?";
        jdbcTemplate.update(sql, title, conversationId);
    }

    // 新增方法：获取会话标题
    public String getConversationTitle(String conversationId) {
        String sql = "SELECT title FROM chat_conversations WHERE id = ?";
        try {
            return jdbcTemplate.queryForObject(sql, String.class, conversationId);
        } catch (Exception e) {
            return null;
        }
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        jdbcTemplate.update("DELETE FROM chat_messages WHERE conversation_id = ?", conversationId);
    }

    @Override
    public List<String> findConversationIds() {
        return jdbcTemplate.queryForList(
                "SELECT DISTINCT conversation_id FROM chat_messages",
                String.class
        );
    }
}
